#include "duckdb/storage/statistics/variant_stats.hpp"
#include "duckdb/storage/statistics/list_stats.hpp"
#include "duckdb/storage/statistics/struct_stats.hpp"
#include "duckdb/storage/statistics/base_statistics.hpp"
#include "duckdb/function/scalar/variant_utils.hpp"

#include "duckdb/common/types/vector.hpp"
#include "duckdb/common/types/variant.hpp"

#include "duckdb/common/serializer/serializer.hpp"
#include "duckdb/common/serializer/deserializer.hpp"

#include "duckdb/common/types/variant_visitor.hpp"
#include "duckdb/function/variant/variant_shredding.hpp"

namespace duckdb {

static void AssertVariant(const BaseStatistics &stats) {
	if (DUCKDB_UNLIKELY(stats.GetStatsType() != StatisticsType::VARIANT_STATS)) {
		throw InternalException(
		    "Calling a VariantStats method on BaseStatistics that are not of type VARIANT, but of type %s",
		    EnumUtil::ToString(stats.GetStatsType()));
	}
}

void VariantStats::Construct(BaseStatistics &stats) {
	stats.child_stats = unsafe_unique_array<BaseStatistics>(new BaseStatistics[2]);
	GetDataUnsafe(stats).shredding_state = VariantStatsShreddingState::UNINITIALIZED;
	CreateUnshreddedStats(stats);
}

BaseStatistics VariantStats::CreateUnknown(LogicalType type) {
	BaseStatistics result(std::move(type));
	result.InitializeUnknown();
	//! Unknown - we have no clue what's in this
	GetDataUnsafe(result).shredding_state = VariantStatsShreddingState::INCONSISTENT;
	result.child_stats[0].Copy(BaseStatistics::CreateUnknown(VariantShredding::GetUnshreddedType()));
	return result;
}

BaseStatistics VariantStats::CreateEmpty(LogicalType type) {
	BaseStatistics result(std::move(type));
	result.InitializeEmpty();
	GetDataUnsafe(result).shredding_state = VariantStatsShreddingState::UNINITIALIZED;
	result.child_stats[0].Copy(BaseStatistics::CreateEmpty(VariantShredding::GetUnshreddedType()));
	return result;
}

//===--------------------------------------------------------------------===//
// Unshredded Stats
//===--------------------------------------------------------------------===//

void VariantStats::CreateUnshreddedStats(BaseStatistics &stats) {
	BaseStatistics::Construct(stats.child_stats[0], VariantShredding::GetUnshreddedType());
}

const BaseStatistics &VariantStats::GetUnshreddedStats(const BaseStatistics &stats) {
	AssertVariant(stats);
	return stats.child_stats[0];
}

BaseStatistics &VariantStats::GetUnshreddedStats(BaseStatistics &stats) {
	AssertVariant(stats);
	return stats.child_stats[0];
}

void VariantStats::SetUnshreddedStats(BaseStatistics &stats, const BaseStatistics &new_stats) {
	AssertVariant(stats);
	stats.child_stats[0].Copy(new_stats);
}

void VariantStats::SetUnshreddedStats(BaseStatistics &stats, unique_ptr<BaseStatistics> new_stats) {
	AssertVariant(stats);
	if (!new_stats) {
		CreateUnshreddedStats(stats);
	} else {
		SetUnshreddedStats(stats, *new_stats);
	}
}

void VariantStats::MarkAsNotShredded(BaseStatistics &stats) {
	D_ASSERT(!IsShredded(stats));
	auto &data = GetDataUnsafe(stats);
	//! All Variant stats start off as UNINITIALIZED, to support merging
	//! This method marks the stats as being unshredded, so they produce INCONSISTENT when merged with SHREDDED stats
	data.shredding_state = VariantStatsShreddingState::NOT_SHREDDED;
}

//===--------------------------------------------------------------------===//
// Shredded Stats
//===--------------------------------------------------------------------===//

static void AssertShreddedStats(const BaseStatistics &stats) {
	if (stats.GetType().id() != LogicalTypeId::STRUCT) {
		throw InternalException("Shredded stats should be of type STRUCT, not %s",
		                        EnumUtil::ToString(stats.GetType().id()));
	}
	auto &struct_children = StructType::GetChildTypes(stats.GetType());
	if (struct_children.size() != 2) {
		throw InternalException(
		    "Shredded stats need to consist of 2 children, 'untyped_value_index' and 'typed_value', not: %s",
		    stats.GetType().ToString());
	}
	if (struct_children[0].second.id() != LogicalTypeId::UINTEGER) {
		throw InternalException("Shredded stats 'untyped_value_index' should be of type UINTEGER, not %s",
		                        EnumUtil::ToString(struct_children[0].second.id()));
	}
}

bool VariantShreddedStats::IsFullyShredded(const BaseStatistics &stats) {
	AssertShreddedStats(stats);

	auto &untyped_value_index_stats = StructStats::GetChildStats(stats, 0);
	auto &typed_value_stats = StructStats::GetChildStats(stats, 1);

	if (!typed_value_stats.CanHaveNull()) {
		//! Fully shredded, no nulls
		return true;
	}
	if (!untyped_value_index_stats.CanHaveNoNull()) {
		//! In the event that this field is entirely missing from the parent OBJECT, both are NULL
		return false;
	}
	if (!NumericStats::HasMin(untyped_value_index_stats) || !NumericStats::HasMax(untyped_value_index_stats)) {
		//! Has no min/max values, essentially double-checking the CanHaveNoNull from above
		return false;
	}
	auto min_value = NumericStats::GetMinUnsafe<uint32_t>(untyped_value_index_stats);
	auto max_value = NumericStats::GetMaxUnsafe<uint32_t>(untyped_value_index_stats);
	if (min_value != max_value) {
		//! Not a constant
		return false;
	}
	//! 0 is reserved for NULL Variant values
	return min_value == 0;
}

LogicalType ToStructuredType(const LogicalType &shredding) {
	D_ASSERT(shredding.id() == LogicalTypeId::STRUCT);
	auto &child_types = StructType::GetChildTypes(shredding);
	D_ASSERT(child_types.size() == 2);

	auto &typed_value = child_types[1].second;

	if (typed_value.id() == LogicalTypeId::STRUCT) {
		auto &struct_children = StructType::GetChildTypes(typed_value);
		child_list_t<LogicalType> structured_children;
		for (auto &child : struct_children) {
			structured_children.emplace_back(child.first, ToStructuredType(child.second));
		}
		return LogicalType::STRUCT(structured_children);
	} else if (typed_value.id() == LogicalTypeId::LIST) {
		auto &child_type = ListType::GetChildType(typed_value);
		return LogicalType::LIST(ToStructuredType(child_type));
	} else {
		return typed_value;
	}
}

LogicalType VariantStats::GetShreddedStructuredType(const BaseStatistics &stats) {
	D_ASSERT(IsShredded(stats));
	return ToStructuredType(GetShreddedStats(stats).GetType());
}

void VariantStats::CreateShreddedStats(BaseStatistics &stats, const LogicalType &shredded_type) {
	BaseStatistics::Construct(stats.child_stats[1], shredded_type);
	auto &data = GetDataUnsafe(stats);
	data.shredding_state = VariantStatsShreddingState::SHREDDED;
}

bool VariantStats::IsShredded(const BaseStatistics &stats) {
	auto &data = GetDataUnsafe(stats);
	return data.shredding_state == VariantStatsShreddingState::SHREDDED;
}

BaseStatistics VariantStats::CreateShredded(const LogicalType &shredded_type) {
	BaseStatistics result(LogicalType::VARIANT());
	result.InitializeEmpty();

	CreateShreddedStats(result, shredded_type);
	result.child_stats[0].Copy(BaseStatistics::CreateEmpty(VariantShredding::GetUnshreddedType()));
	result.child_stats[1].Copy(BaseStatistics::CreateEmpty(shredded_type));
	return result;
}

const BaseStatistics &VariantStats::GetShreddedStats(const BaseStatistics &stats) {
	AssertVariant(stats);
	D_ASSERT(IsShredded(stats));
	return stats.child_stats[1];
}

BaseStatistics &VariantStats::GetShreddedStats(BaseStatistics &stats) {
	AssertVariant(stats);
	D_ASSERT(IsShredded(stats));
	return stats.child_stats[1];
}

void VariantStats::SetShreddedStats(BaseStatistics &stats, const BaseStatistics &new_stats) {
	auto &data = GetDataUnsafe(stats);
	if (!IsShredded(stats)) {
		BaseStatistics::Construct(stats.child_stats[1], new_stats.GetType());
		D_ASSERT(data.shredding_state != VariantStatsShreddingState::INCONSISTENT);
		data.shredding_state = VariantStatsShreddingState::SHREDDED;
	}
	stats.child_stats[1].Copy(new_stats);
}

void VariantStats::SetShreddedStats(BaseStatistics &stats, unique_ptr<BaseStatistics> new_stats) {
	AssertVariant(stats);
	D_ASSERT(new_stats);
	SetShreddedStats(stats, *new_stats);
}

//===--------------------------------------------------------------------===//
// (De)Serialization
//===--------------------------------------------------------------------===//

void VariantStats::Serialize(const BaseStatistics &stats, Serializer &serializer) {
	auto &data = GetDataUnsafe(stats);
	auto &unshredded_stats = VariantStats::GetUnshreddedStats(stats);

	serializer.WriteProperty(200, "shredding_state", data.shredding_state);

	serializer.WriteProperty(225, "unshredded_stats", unshredded_stats);
	if (IsShredded(stats)) {
		auto &shredded_stats = VariantStats::GetShreddedStats(stats);
		serializer.WriteProperty(230, "shredded_type", shredded_stats.type);
		serializer.WriteProperty(235, "shredded_stats", shredded_stats);
	}
}

void VariantStats::Deserialize(Deserializer &deserializer, BaseStatistics &base) {
	auto &type = base.GetType();
	D_ASSERT(type.InternalType() == PhysicalType::STRUCT);
	D_ASSERT(type.id() == LogicalTypeId::VARIANT);
	auto &data = GetDataUnsafe(base);

	auto unshredded_type = VariantShredding::GetUnshreddedType();
	data.shredding_state = deserializer.ReadProperty<VariantStatsShreddingState>(200, "shredding_state");

	{
		//! Read the 'unshredded_stats' child
		deserializer.Set<const LogicalType &>(unshredded_type);
		auto stat = deserializer.ReadProperty<BaseStatistics>(225, "unshredded_stats");
		base.child_stats[0].Copy(stat);
		deserializer.Unset<LogicalType>();
	}

	if (!IsShredded(base)) {
		return;
	}
	//! Read the type of the 'shredded_stats'
	auto shredded_type = deserializer.ReadProperty<LogicalType>(230, "shredded_type");

	{
		//! Finally read the 'shredded_stats' themselves
		deserializer.Set<const LogicalType &>(shredded_type);
		auto stat = deserializer.ReadProperty<BaseStatistics>(235, "shredded_stats");
		if (base.child_stats[1].type.id() == LogicalTypeId::INVALID) {
			base.child_stats[1] = BaseStatistics::CreateUnknown(shredded_type);
		}
		base.child_stats[1].Copy(stat);
		deserializer.Unset<LogicalType>();
	}
}

static string ToStringInternal(const BaseStatistics &stats) {
	string result;
	result = StringUtil::Format("fully_shredded: %s", VariantShreddedStats::IsFullyShredded(stats) ? "true" : "false");

	auto &typed_value = StructStats::GetChildStats(stats, 1);
	auto type_id = typed_value.GetType().id();
	if (type_id == LogicalTypeId::LIST) {
		result += ", child: ";
		auto &child_stats = ListStats::GetChildStats(typed_value);
		result += ToStringInternal(child_stats);
	} else if (type_id == LogicalTypeId::STRUCT) {
		result += ", children: {";
		auto &fields = StructType::GetChildTypes(typed_value.GetType());
		for (idx_t i = 0; i < fields.size(); i++) {
			if (i) {
				result += ", ";
			}
			auto &child_stats = StructStats::GetChildStats(typed_value, i);
			result += StringUtil::Format("%s: %s", fields[i].first, ToStringInternal(child_stats));
		}
		result += "}";
	}
	return result;
}

string VariantStats::ToString(const BaseStatistics &stats) {
	string result;
	bool is_shredded = IsShredded(stats);
	auto &data = GetDataUnsafe(stats);
	result = StringUtil::Format("shredding_state: %s", EnumUtil::ToString(data.shredding_state));
	if (is_shredded) {
		result += ", shredding: {";
		result += StringUtil::Format("typed_value_type: %s, ", ToStructuredType(stats.child_stats[1].type).ToString());
		result += StringUtil::Format("stats: {%s}", ToStringInternal(stats.child_stats[1]));
		result += "}";
	}
	return result;
}

static BaseStatistics WrapTypedValue(BaseStatistics &untyped_value_index, BaseStatistics &typed_value) {
	BaseStatistics shredded = BaseStatistics::CreateEmpty(LogicalType::STRUCT(
	    {{"untyped_value_index", untyped_value_index.GetType()}, {"typed_value", typed_value.GetType()}}));

	StructStats::GetChildStats(shredded, 0).Copy(untyped_value_index);
	StructStats::GetChildStats(shredded, 1).Copy(typed_value);
	return shredded;
}

bool VariantStats::MergeShredding(BaseStatistics &stats, const BaseStatistics &other, BaseStatistics &new_stats) {
	//! shredded_type:
	//! STRUCT(untyped_value_index UINTEGER, typed_value <shredding>)

	//! shredding, 1 of:
	//! - <primitive type>
	//! - <shredded_type>
	//! - <shredded_type>[]

	D_ASSERT(stats.type.id() == LogicalTypeId::STRUCT);
	D_ASSERT(other.type.id() == LogicalTypeId::STRUCT);

	auto &stats_children = StructType::GetChildTypes(stats.type);
	auto &other_children = StructType::GetChildTypes(other.type);
	D_ASSERT(stats_children.size() == 2);
	D_ASSERT(other_children.size() == 2);

	auto &stats_typed_value_type = stats_children[1].second;
	auto &other_typed_value_type = other_children[1].second;

	//! Merge the untyped_value_index stats
	auto &untyped_value_index = StructStats::GetChildStats(stats, 0);
	untyped_value_index.Merge(StructStats::GetChildStats(other, 0));

	auto &stats_typed_value = StructStats::GetChildStats(stats, 1);
	auto &other_typed_value = StructStats::GetChildStats(other, 1);

	if (stats_typed_value_type.id() == LogicalTypeId::STRUCT) {
		if (stats_typed_value_type.id() != other_typed_value_type.id()) {
			//! other is not an OBJECT, can't merge
			return false;
		}
		auto &stats_object_children = StructType::GetChildTypes(stats_typed_value_type);
		auto &other_object_children = StructType::GetChildTypes(other_typed_value_type);

		//! Map field name to index, for 'other'
		case_insensitive_map_t<idx_t> key_to_index;
		for (idx_t i = 0; i < other_object_children.size(); i++) {
			auto &other_object_child = other_object_children[i];
			key_to_index.emplace(other_object_child.first, i);
		}

		//! Attempt to merge all overlapping fields, only keep the fields that were able to be merged
		child_list_t<LogicalType> new_children;
		vector<BaseStatistics> new_child_stats;

		for (idx_t i = 0; i < stats_object_children.size(); i++) {
			auto &stats_object_child = stats_object_children[i];
			auto other_it = key_to_index.find(stats_object_child.first);
			if (other_it == key_to_index.end()) {
				continue;
			}
			auto &other_object_child = other_object_children[other_it->second];
			if (other_object_child.second.id() != stats_object_child.second.id()) {
				//! TODO: perhaps we can keep the field but demote the type to unshredded somehow?
				//! Or even use MaxLogicalType and merge the stats into that ?
				continue;
			}

			auto &stats_child = StructStats::GetChildStats(stats_typed_value, i);
			auto &other_child = StructStats::GetChildStats(other_typed_value, other_it->second);
			BaseStatistics new_child;
			if (!MergeShredding(stats_child, other_child, new_child)) {
				continue;
			}
			new_children.emplace_back(stats_object_child.first, new_child.GetType());
			new_child_stats.emplace_back(std::move(new_child));
		}
		if (new_children.empty()) {
			//! No fields remaining, demote to unshredded
			return false;
		}

		//! Create new stats out of the remaining fields
		auto new_object_type = LogicalType::STRUCT(std::move(new_children));
		auto new_typed_value = BaseStatistics::CreateEmpty(new_object_type);
		for (idx_t i = 0; i < new_child_stats.size(); i++) {
			StructStats::SetChildStats(new_typed_value, i, new_child_stats[i]);
		}
		new_typed_value.CombineValidity(stats_typed_value, other_typed_value);
		new_stats = WrapTypedValue(untyped_value_index, new_typed_value);
		return true;
	} else if (stats_typed_value_type.id() == LogicalTypeId::LIST) {
		if (stats_typed_value_type.id() != other_typed_value_type.id()) {
			//! other is not an ARRAY, can't merge
			return false;
		}
		auto &stats_child = ListStats::GetChildStats(stats_typed_value);
		auto &other_child = ListStats::GetChildStats(other_typed_value);

		//! TODO: perhaps we can keep the LIST part of the stats, and only demote the child to unshredded?
		BaseStatistics new_child_stats;
		if (!MergeShredding(stats_child, other_child, new_child_stats)) {
			return false;
		}
		auto new_typed_value = BaseStatistics::CreateEmpty(LogicalType::LIST(new_child_stats.type));
		new_typed_value.CombineValidity(stats_typed_value, other_typed_value);
		ListStats::SetChildStats(new_typed_value, new_child_stats.ToUnique());
		new_stats = WrapTypedValue(untyped_value_index, new_typed_value);
		return true;
	} else {
		D_ASSERT(!stats_typed_value_type.IsNested());
		if (stats_typed_value_type.id() != other_typed_value_type.id()) {
			//! other is not the same type, can't merge
			return false;
		}
		stats_typed_value.Merge(other_typed_value);
		new_stats = std::move(stats);
		return true;
	}
}

void VariantStats::Merge(BaseStatistics &stats, const BaseStatistics &other) {
	if (other.GetType().id() == LogicalTypeId::VALIDITY) {
		return;
	}

	stats.child_stats[0].Merge(other.child_stats[0]);
	auto &data = GetDataUnsafe(stats);
	auto &other_data = GetDataUnsafe(other);

	const auto other_shredding_state = other_data.shredding_state;
	const auto shredding_state = data.shredding_state;

	if (other_shredding_state == VariantStatsShreddingState::UNINITIALIZED) {
		//! No need to merge
		return;
	}

	switch (shredding_state) {
	case VariantStatsShreddingState::INCONSISTENT: {
		//! INCONSISTENT + ANY -> INCONSISTENT
		return;
	}
	case VariantStatsShreddingState::UNINITIALIZED: {
		switch (other_shredding_state) {
		case VariantStatsShreddingState::SHREDDED:
			stats.child_stats[1] = BaseStatistics::CreateUnknown(other.child_stats[1].GetType());
			stats.child_stats[1].Copy(other.child_stats[1]);
			break;
		default:
			break;
		}
		//! UNINITIALIZED + ANY -> ANY
		data.shredding_state = other_shredding_state;
		break;
	}
	case VariantStatsShreddingState::NOT_SHREDDED: {
		if (other_shredding_state == VariantStatsShreddingState::NOT_SHREDDED) {
			return;
		}
		//! NOT_SHREDDED + !NOT_SHREDDED -> INCONSISTENT
		data.shredding_state = VariantStatsShreddingState::INCONSISTENT;
		stats.child_stats[1].type = LogicalType::INVALID;
		break;
	}
	case VariantStatsShreddingState::SHREDDED: {
		switch (other_shredding_state) {
		case VariantStatsShreddingState::SHREDDED: {
			BaseStatistics merged_shredding_stats;
			if (!MergeShredding(stats.child_stats[1], other.child_stats[1], merged_shredding_stats)) {
				//! SHREDDED(T1) + SHREDDED(T2) -> INCONSISTENT
				data.shredding_state = VariantStatsShreddingState::INCONSISTENT;
				stats.child_stats[1].type = LogicalType::INVALID;
			} else {
				//! SHREDDED(T1) + SHREDDED(T1) -> SHREDDED
				stats.child_stats[1] = BaseStatistics::CreateUnknown(merged_shredding_stats.GetType());
				stats.child_stats[1].Copy(merged_shredding_stats);
			}
			break;
		}
		default:
			//! SHREDDED + !SHREDDED -> INCONSISTENT
			data.shredding_state = VariantStatsShreddingState::INCONSISTENT;
			stats.child_stats[1].type = LogicalType::INVALID;
			break;
		}
		break;
	}
	}
}

void VariantStats::Copy(BaseStatistics &stats, const BaseStatistics &other) {
	auto &other_data = VariantStats::GetDataUnsafe(other);
	auto &data = VariantStats::GetDataUnsafe(stats);
	(void)data;

	//! This is ensured by the CopyBase method of BaseStatistics
	D_ASSERT(data.shredding_state == other_data.shredding_state);
	stats.child_stats[0].Copy(other.child_stats[0]);
	if (IsShredded(other)) {
		stats.child_stats[1] = BaseStatistics::CreateUnknown(other.child_stats[1].GetType());
		stats.child_stats[1].Copy(other.child_stats[1]);
	} else {
		stats.child_stats[1].type = LogicalType::INVALID;
	}
}

void VariantStats::Verify(const BaseStatistics &stats, Vector &vector, const SelectionVector &sel, idx_t count) {
	// TODO: Verify stats
}

const VariantStatsData &VariantStats::GetDataUnsafe(const BaseStatistics &stats) {
	AssertVariant(stats);
	return stats.stats_union.variant_data;
}

VariantStatsData &VariantStats::GetDataUnsafe(BaseStatistics &stats) {
	AssertVariant(stats);
	return stats.stats_union.variant_data;
}

} // namespace duckdb
